import cv2
import datetime
videoSrc = "../vtest.avi"

# 参数为0是打开摄像头，文件名是打开视频
cap = cv2.VideoCapture(videoSrc)
# 混合高斯背景建模算法
fgbg = cv2.createBackgroundSubtractorMOG2()
# 指定写视频的格式
fourcc = cv2.VideoWriter_fourcc(*'XVID')
formateTime = datetime.datetime.now().strftime("%A_%d_%B_%Y_%I_%M_%S%p") + '.avi'

# 获得码率
fps = cap.get(cv2.CAP_PROP_FPS)
print(fps)
# 获取视频的尺寸
size = (int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)),
        int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)))
print(size)
out = cv2.VideoWriter(formateTime, fourcc, 10, fps, size)

while True:
    ret, frame = cap.read()
    fgmask = fgbg.apply(frame)

    # 形态学去噪
    element = cv2.getStructuringElement(cv2.MORPH_CROSS, (3, 3))
    # 开运算去噪
    cv2.morphologyEx(fgmask, cv2.MORPH_OPEN, element)
    # 寻找角点
    _, contours, hierarchy = cv2.findContours(fgmask.copy(), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    count = 0
    for cont in contours:
        area = cv2.contourArea(cont)
        if area < 300:
            continue
        # 数量加1
        count += 1

        # 提取矩形坐标及宽、高
        rect = cv2.boundingRect(cont)
        (x, y, w, h) = rect


        # 原图上绘制矩形
        cv2.rectangle(frame, (x, y), (x + w, y + h), (0, 205, 205), 1)
        # 黑白前景上绘制矩形
        cv2.rectangle(fgmask, (x, y), (x + w, y + h), (0xff, 0xff, 0xff), 1)

        # 防止编号到图片之外
        yy = 10 if y < 10 else y
        cv2.putText(frame, str(count), (x, yy), cv2.FONT_HERSHEY_COMPLEX, 0.4, (0, 255, 0), 1)

    # 显示总数
    cv2.putText(frame, "count", (5, 20), cv2.FONT_HERSHEY_COMPLEX, 0.6, (0, 255, 0), 1)
    cv2.putText(frame, str(count), (5, 20), cv2.FONT_HERSHEY_COMPLEX, 0.6, (0, 255, 0), 1)
    print("***********************")

    # 原图上标注
    cv2.imshow("frame", frame)
    # 以黑白的形式显示前景和背景
    cv2.imshow("frame2", fgmask)
    #
    out.write(frame)
    # 按esc退出
    k = cv2.waitKey(30) & 0xff
    if k == 27:
        break
out.release()
cap.release()
cv2.destroyAllWindows()
